import torch
import torch.nn as nn
from torch.nn import BCEWithLogitsLoss
from segmentation_models_pytorch.losses.dice import DiceLoss
class CombinedLoss(nn.Module):
    def __init__(self, mode ='binary',dice_weight=0.5, ce_weight=0.5):
        super(CombinedLoss, self).__init__()
        self.dice_loss = DiceLoss(mode=mode)
        self.ce_loss = BCEWithLogitsLoss()
        self.dice_weight = dice_weight
        self.ce_weight = ce_weight

    def forward(self, inputs, targets):
        dice_loss = self.dice_loss(inputs, targets)
        ce_loss = self.ce_loss(inputs, targets)
        return self.dice_weight * dice_loss + self.ce_weight * ce_loss
